{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.0.0-beta1\n",
      "sys.version_info(major=3, minor=5, micro=3, releaselevel='final', serial=0)\n",
      "matplotlib 3.0.3\n",
      "numpy 1.16.4\n",
      "pandas 0.24.2\n",
      "sklearn 0.21.2\n",
      "tensorflow 2.0.0-beta1\n",
      "tensorflow.python.keras.api._v2.keras 2.2.4-tf\n"
     ]
    }
   ],
   "source": [
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import sklearn\n",
    "import pandas as pd\n",
    "import os\n",
    "import sys\n",
    "import time\n",
    "import tensorflow as tf\n",
    "\n",
    "from tensorflow import keras\n",
    "\n",
    "print(tf.__version__)\n",
    "print(sys.version_info)\n",
    "for module in mpl, np, pd, sklearn, tf, keras:\n",
    "    print(module.__name__, module.__version__)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4\n",
      "4\n"
     ]
    }
   ],
   "source": [
    "tf.debugging.set_log_device_placement(True)\n",
    "gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "for gpu in gpus:\n",
    "    tf.config.experimental.set_memory_growth(gpu, True)\n",
    "print(len(gpus))\n",
    "logical_gpus = tf.config.experimental.list_logical_devices('GPU')\n",
    "print(len(logical_gpus))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(5000, 28, 28) (5000,)\n",
      "(55000, 28, 28) (55000,)\n",
      "(10000, 28, 28) (10000,)\n"
     ]
    }
   ],
   "source": [
    "fashion_mnist = keras.datasets.fashion_mnist\n",
    "(x_train_all, y_train_all), (x_test, y_test) = fashion_mnist.load_data()\n",
    "x_valid, x_train = x_train_all[:5000], x_train_all[5000:]\n",
    "y_valid, y_train = y_train_all[:5000], y_train_all[5000:]\n",
    "\n",
    "print(x_valid.shape, y_valid.shape)\n",
    "print(x_train.shape, y_train.shape)\n",
    "print(x_test.shape, y_test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "scaler = StandardScaler()\n",
    "x_train_scaled = scaler.fit_transform(\n",
    "    x_train.astype(np.float32).reshape(-1, 1)).reshape(-1, 28, 28, 1)\n",
    "x_valid_scaled = scaler.transform(\n",
    "    x_valid.astype(np.float32).reshape(-1, 1)).reshape(-1, 28, 28, 1)\n",
    "x_test_scaled = scaler.transform(\n",
    "    x_test.astype(np.float32).reshape(-1, 1)).reshape(-1, 28, 28, 1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Executing op TensorSliceDataset in device /job:localhost/replica:0/task:0/device:CPU:0\n",
      "Executing op ShuffleDataset in device /job:localhost/replica:0/task:0/device:CPU:0\n",
      "Executing op RepeatDataset in device /job:localhost/replica:0/task:0/device:CPU:0\n",
      "Executing op BatchDatasetV2 in device /job:localhost/replica:0/task:0/device:CPU:0\n",
      "Executing op PrefetchDataset in device /job:localhost/replica:0/task:0/device:CPU:0\n",
      "Executing op ExperimentalRebatchDataset in device /job:localhost/replica:0/task:0/device:CPU:0\n",
      "Executing op ExperimentalAutoShardDataset in device /job:localhost/replica:0/task:0/device:CPU:0\n"
     ]
    }
   ],
   "source": [
    "def make_dataset(images, labels, epochs, batch_size, shuffle=True):\n",
    "    dataset = tf.data.Dataset.from_tensor_slices((images, labels))\n",
    "    if shuffle:\n",
    "        dataset = dataset.shuffle(10000)\n",
    "    dataset = dataset.repeat(epochs).batch(batch_size).prefetch(50)\n",
    "    return dataset\n",
    "\n",
    "strategy = tf.distribute.MirroredStrategy()\n",
    "\n",
    "with strategy.scope():\n",
    "    batch_size_per_replica = 256\n",
    "    batch_size = batch_size_per_replica * len(logical_gpus)\n",
    "    train_dataset = make_dataset(x_train_scaled, y_train, 1, batch_size)\n",
    "    valid_dataset = make_dataset(x_valid_scaled, y_valid, 1, batch_size)\n",
    "    train_dataset_distribute = strategy.experimental_distribute_dataset(\n",
    "        train_dataset)\n",
    "    valid_dataset_distribute = strategy.experimental_distribute_dataset(\n",
    "        valid_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Executing op RandomUniform in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op Sub in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op Mul in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op Add in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op VarIsInitializedOp in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op LogicalNot in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op Assert in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op VarIsInitializedOp in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op LogicalNot in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op Assert in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op VarIsInitializedOp in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op LogicalNot in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op Assert in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:3\n",
      "Executing op VarIsInitializedOp in device /job:localhost/replica:0/task:0/device:GPU:3\n",
      "Executing op LogicalNot in device /job:localhost/replica:0/task:0/device:GPU:3\n",
      "Executing op Assert in device /job:localhost/replica:0/task:0/device:GPU:3\n",
      "Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:3\n",
      "Executing op Fill in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:3\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:3\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:3\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:3\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:3\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:3\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:3\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:3\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:3\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:3\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:3\n"
     ]
    }
   ],
   "source": [
    "with strategy.scope():\n",
    "    model = keras.models.Sequential()\n",
    "    model.add(keras.layers.Conv2D(filters=128, kernel_size=3,\n",
    "                                  padding='same',\n",
    "                                  activation='relu',\n",
    "                                  input_shape=(28, 28, 1)))\n",
    "    model.add(keras.layers.Conv2D(filters=128, kernel_size=3,\n",
    "                                  padding='same',\n",
    "                                  activation='relu'))\n",
    "    model.add(keras.layers.MaxPool2D(pool_size=2))\n",
    "    model.add(keras.layers.Conv2D(filters=256, kernel_size=3,\n",
    "                                  padding='same',\n",
    "                                  activation='relu'))\n",
    "    model.add(keras.layers.Conv2D(filters=256, kernel_size=3,\n",
    "                                  padding='same',\n",
    "                                  activation='relu'))\n",
    "    model.add(keras.layers.MaxPool2D(pool_size=2))\n",
    "    model.add(keras.layers.Conv2D(filters=512, kernel_size=3,\n",
    "                                  padding='same',\n",
    "                                  activation='relu'))\n",
    "    model.add(keras.layers.Conv2D(filters=512, kernel_size=3,\n",
    "                                  padding='same',\n",
    "                                  activation='relu'))\n",
    "    model.add(keras.layers.MaxPool2D(pool_size=2))\n",
    "    model.add(keras.layers.Flatten())\n",
    "    model.add(keras.layers.Dense(512, activation='relu'))\n",
    "    model.add(keras.layers.Dense(10, activation=\"softmax\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "conv2d (Conv2D)              (None, 28, 28, 128)       1280      \n",
      "_________________________________________________________________\n",
      "conv2d_1 (Conv2D)            (None, 28, 28, 128)       147584    \n",
      "_________________________________________________________________\n",
      "max_pooling2d (MaxPooling2D) (None, 14, 14, 128)       0         \n",
      "_________________________________________________________________\n",
      "conv2d_2 (Conv2D)            (None, 14, 14, 256)       295168    \n",
      "_________________________________________________________________\n",
      "conv2d_3 (Conv2D)            (None, 14, 14, 256)       590080    \n",
      "_________________________________________________________________\n",
      "max_pooling2d_1 (MaxPooling2 (None, 7, 7, 256)         0         \n",
      "_________________________________________________________________\n",
      "conv2d_4 (Conv2D)            (None, 7, 7, 512)         1180160   \n",
      "_________________________________________________________________\n",
      "conv2d_5 (Conv2D)            (None, 7, 7, 512)         2359808   \n",
      "_________________________________________________________________\n",
      "max_pooling2d_2 (MaxPooling2 (None, 3, 3, 512)         0         \n",
      "_________________________________________________________________\n",
      "flatten (Flatten)            (None, 4608)              0         \n",
      "_________________________________________________________________\n",
      "dense (Dense)                (None, 512)               2359808   \n",
      "_________________________________________________________________\n",
      "dense_1 (Dense)              (None, 10)                5130      \n",
      "=================================================================\n",
      "Total params: 6,939,018\n",
      "Trainable params: 6,939,018\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:3\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:3\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:3\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:2\n",
      "Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:3\n",
      "Executing op __inference_initialize_variables_10549 in device <unspecified>\n",
      "Executing op __inference_distributed_train_step_11728 in device <unspecified>\n",
      "total: 107.368, num_batches: 53, average: 2.026, time: 0.203Executing op __inference_distributed_train_step_13332 in device <unspecified>\n",
      "total: 108.739, num_batches: 54, average: 2.014, time: 37.324Executing op __inference_distributed_test_step_13736 in device <unspecified>\n",
      "Executing op __inference_distributed_test_step_14132 in device <unspecified>\n",
      "Epoch: 1, Loss: 2.014, Acc: 0.810, Val_Loss: 0.440, Val_Acc: 0.846\n",
      "Epoch: 2, Loss: 1.872, Acc: 0.823, Val_Loss: 0.411, Val_Acc: 0.857\n",
      "Epoch: 3, Loss: 1.750, Acc: 0.837, Val_Loss: 0.383, Val_Acc: 0.864\n",
      "Epoch: 4, Loss: 1.664, Acc: 0.845, Val_Loss: 0.370, Val_Acc: 0.867\n",
      "Epoch: 5, Loss: 1.585, Acc: 0.852, Val_Loss: 0.357, Val_Acc: 0.870\n",
      "Epoch: 6, Loss: 1.519, Acc: 0.858, Val_Loss: 0.345, Val_Acc: 0.875\n",
      "Epoch: 7, Loss: 1.467, Acc: 0.863, Val_Loss: 0.335, Val_Acc: 0.879\n",
      "Epoch: 8, Loss: 1.409, Acc: 0.869, Val_Loss: 0.325, Val_Acc: 0.883\n",
      "Epoch: 9, Loss: 1.424, Acc: 0.869, Val_Loss: 0.326, Val_Acc: 0.880\n",
      "Epoch: 10, Loss: 1.346, Acc: 0.875, Val_Loss: 0.317, Val_Acc: 0.884\n"
     ]
    }
   ],
   "source": [
    "# customized training loop.\n",
    "# 1. define losses functions\n",
    "# 2. define function train_step\n",
    "# 3. define function test_step\n",
    "# 4. for-loop training loop\n",
    "\n",
    "with strategy.scope():\n",
    "    # batch_size, batch_size / #{gpu}\n",
    "    # eg: 64, gpu: 16\n",
    "    loss_func = keras.losses.SparseCategoricalCrossentropy(\n",
    "        reduction = keras.losses.Reduction.NONE)\n",
    "    def compute_loss(labels, predictions):\n",
    "        per_replica_loss = loss_func(labels, predictions)\n",
    "        return tf.nn.compute_average_loss(per_replica_loss,\n",
    "                                          global_batch_size = batch_size)\n",
    "    \n",
    "    test_loss = keras.metrics.Mean(name = \"test_loss\")\n",
    "    train_accuracy = keras.metrics.SparseCategoricalAccuracy(\n",
    "        name = 'train_accuracy')\n",
    "    test_accuracy = keras.metrics.SparseCategoricalAccuracy(\n",
    "        name = 'test_accuracy')\n",
    "\n",
    "    optimizer = keras.optimizers.SGD(lr=0.01)\n",
    "\n",
    "    def train_step(inputs):\n",
    "        images, labels = inputs\n",
    "        with tf.GradientTape() as tape:\n",
    "            predictions = model(images, training = True)\n",
    "            loss = compute_loss(labels, predictions)\n",
    "        gradients = tape.gradient(loss, model.trainable_variables)\n",
    "        optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "        train_accuracy.update_state(labels, predictions)\n",
    "        return loss\n",
    "    \n",
    "    @tf.function\n",
    "    def distributed_train_step(inputs):\n",
    "        per_replica_average_loss = strategy.experimental_run_v2(\n",
    "            train_step, args = (inputs,))\n",
    "        return strategy.reduce(tf.distribute.ReduceOp.SUM,\n",
    "                               per_replica_average_loss,\n",
    "                               axis = None)\n",
    "    \n",
    "    def test_step(inputs):\n",
    "        images, labels = inputs\n",
    "        predictions = model(images)\n",
    "        t_loss = loss_func(labels, predictions)\n",
    "        test_loss.update_state(t_loss)\n",
    "        test_accuracy.update_state(labels, predictions)\n",
    "        \n",
    "    @tf.function\n",
    "    def distributed_test_step(inputs):\n",
    "        strategy.experimental_run_v2(\n",
    "            test_step, args = (inputs,))\n",
    "\n",
    "    epochs = 10\n",
    "    for epoch in range(epochs):\n",
    "        total_loss = 0.0\n",
    "        num_batches = 0\n",
    "        for x in train_dataset:\n",
    "            start_time = time.time()\n",
    "            total_loss += distributed_train_step(x)\n",
    "            run_time = time.time() - start_time\n",
    "            num_batches += 1\n",
    "            print('\\rtotal: %3.3f, num_batches: %d, '\n",
    "                  'average: %3.3f, time: %3.3f'\n",
    "                  % (total_loss, num_batches,\n",
    "                     total_loss / num_batches, run_time),\n",
    "                  end = '')\n",
    "        train_loss = total_loss / num_batches\n",
    "        for x in valid_dataset:\n",
    "            distributed_test_step(x)\n",
    "\n",
    "        print('\\rEpoch: %d, Loss: %3.3f, Acc: %3.3f, '\n",
    "              'Val_Loss: %3.3f, Val_Acc: %3.3f'\n",
    "              % (epoch + 1, train_loss, train_accuracy.result(),\n",
    "                 test_loss.result(), test_accuracy.result()))\n",
    "        test_loss.reset_states()\n",
    "        train_accuracy.reset_states()\n",
    "        test_accuracy.reset_states()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
